import os

from model.decoder.decoder import Decoder

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC']
plt.rcParams['axes.unicode_minus'] = False

import torch
from torchvision import transforms
from PIL import Image
import numpy as np


from model.encoder.encoder import Encoder
from options import HiDDenConfiguration

def load_image(image_path, target_size=(64, 64)):
    """加载并预处理图像，转换为模型输入格式的张量"""
    try:
        # 打开图像并转换为 RGB 格式（排除 RGBA 或灰度图）
        img = Image.open(image_path).convert('RGB')
        # 定义预处理流程：调整尺寸 → 转为张量
        preprocess = transforms.Compose([
            transforms.Resize(target_size),
            transforms.ToTensor(),  # 输出范围 [0, 1]，形状 [3, H, W]
        ])
        # 预处理并添加批次维度
        img_tensor = preprocess(img).unsqueeze(0)  # 形状 [1, 3, H, W]
        return img_tensor
    except Exception as e:
        raise ValueError(f"图像加载失败：{str(e)}")

# 配置参数
config = HiDDenConfiguration(
    H=640, W=427,
    message_length=100,
    encoder_blocks=4, encoder_channels=64,
    decoder_blocks=4, decoder_channels=64,
    use_discriminator=True, use_vgg=True,
    discriminator_blocks=4, discriminator_channels=64,
    decoder_loss=1.0, encoder_loss=1.0, adversarial_loss=0.1
)


image_path = "test_image.jpg"
# 加载单张图像并扩展为批次（batch_size=2）
image = load_image(image_path, target_size=(config.H, config.W))
batch_size = 2
image = image.repeat(batch_size, 1, 1, 1)  # 重复生成批次数据 [2, 3, 64, 64]

# 生成随机秘密信息（二进制）
message = torch.randint(0, 2, (batch_size, config.message_length)).float()

# 初始化编码器并生成水印图像
encoder = Encoder(config)
decoder = Decoder(config)
with torch.no_grad():  # 关闭梯度计算，加速推理
    watermarked_image = encoder(image, message)
    extracted_message = decoder(watermarked_image)
    extracted_message = (extracted_message > 0.5).float()

# 验证输出尺寸
assert watermarked_image.shape == image.shape, \
    f"输出尺寸错误！预期 {image.shape}，实际 {watermarked_image.shape}"
print("输出尺寸验证通过！")

# 打印第0个样本的原始信息和提取信息
print("原始秘密信息（前20位）：", message[0, :20].numpy())
print("提取的秘密信息（前20位）：", extracted_message[0, :20].numpy())

# 计算每个样本的准确率，再求平均值
accuracy = (extracted_message == message).float().mean(dim=1).mean()
print(f"秘密信息提取准确率：{accuracy.item():.4f}")  # 理想情况下应接近1.0


def tensor_to_image(tensor):
    """
    将张量转换为可显示的图像格式（兼容单张/批次图像，自动适配范围）
    Args:
        tensor: 输入张量，形状可为 [C, H, W]（单张）或 [batch, C, H, W]（批次）
    Returns:
        图像数组，形状 [H, W, C] 或 [batch, H, W, C]，dtype=uint8
    """
    # 移到CPU并转为numpy（若已是numpy则跳过）
    if isinstance(tensor, torch.Tensor):
        tensor = tensor.cpu().detach().numpy()

    # 处理维度：若为单张图像（无batch维），先增加batch维再转换
    has_batch_dim = tensor.ndim == 4
    if not has_batch_dim:
        tensor = tensor[np.newaxis, ...]  # [C, H, W] → [1, C, H, W]

    # 转换维度为 [batch, H, W, C]
    image = tensor.transpose(0, 2, 3, 1)  # 等效于permute(0,2,3,1)，但numpy用transpose

    # 自动判断数值范围并映射到[0,255]
    if image.min() < 0:  # 假设范围是[-1, 1]
        image = (image + 1) * 127.5
    else:  # 假设范围是[0, 1]
        image = image * 255

    # 截断并转换为uint8
    image = np.clip(image, 0, 255).astype(np.uint8)

    # 若输入是单张图像，去除batch维
    if not has_batch_dim:
        image = image[0]

    return image


# 可视化对比（显示批次中第一张图像）
plt.figure(figsize=(10, 5))

# 原始图像
plt.subplot(121)
plt.title("原始图像")
plt.imshow(tensor_to_image(image[0]))
plt.axis('off')  # 关闭坐标轴

# 含水印图像
plt.subplot(122)
plt.title("含水印图像")
plt.imshow(tensor_to_image(watermarked_image[0]))
plt.axis('off')

plt.tight_layout()  # 自动调整子图间距
plt.show()